# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# python /data0/liqifeng/ZYC/vggt/demo_gradio.py --ckpt /data0/liqifeng/ZYC/model.pt
# 直接读取预计算的特征文件并只跑 heads（跳过 backbone/aggregator）

import os
import json
import cv2
import torch
import numpy as np
import gradio as gr
import sys
import shutil
from datetime import datetime
import glob
import gc
import time
import argparse
from importlib import import_module
from importlib import util as importlib_util

sys.path.append("vggt/")

from visual_util import predictions_to_glb
from vggt.models.vggt import VGGT
from vggt.utils.load_fn import load_and_preprocess_images
from vggt.utils.pose_enc import pose_encoding_to_extri_intri
from vggt.utils.geometry import unproject_depth_map_to_point_map

h5py_spec = importlib_util.find_spec("h5py")
if h5py_spec is not None:  # pragma: no cover - dependency import is environment-specific
    h5py = import_module("h5py")
else:  # pragma: no cover - dependency import is environment-specific
    h5py = None

device = "cuda" if torch.cuda.is_available() else "cpu"

print("Initializing and loading VGGT model...")
# model = VGGT.from_pretrained("facebook/VGGT-1B")  # another way to load the model


# -------------------------------------------------------------------------
# Utilities: export point cloud to PLY/NPZ
# -------------------------------------------------------------------------
def _parse_frame_filter(frame_filter):
    """Return indices to keep based on the frame_filter dropdown value."""
    if frame_filter in (None, "All"):
        return slice(None)
    try:
        idx = int(str(frame_filter).split(":")[0].strip())
        return [idx]
    except Exception:
        return slice(None)


def _frame_filter_to_index(frame_filter, num_frames: int) -> int:
    """Resolve a single frame index from the dropdown selection."""
    if isinstance(frame_filter, (list, tuple)) and frame_filter:
        candidate = frame_filter[0]
        if isinstance(candidate, int):
            return max(0, min(num_frames - 1, candidate))
    if isinstance(frame_filter, str) and frame_filter not in ("", "All"):
        try:
            raw_idx = int(frame_filter.split(":")[0].strip())
            return max(0, min(num_frames - 1, raw_idx))
        except Exception:
            pass
    return 0


def _render_depth_preview(depth_array: np.ndarray, frame_filter) -> np.ndarray | None:
    """Create a colorized preview image for the selected depth frame."""
    if depth_array is None:
        return None

    depth_np = np.asarray(depth_array)
    if depth_np.ndim == 4 and depth_np.shape[-1] == 1:
        depth_np = depth_np[..., 0]
    if depth_np.ndim != 3:
        return None

    num_frames = depth_np.shape[0]
    if num_frames == 0:
        return None

    frame_idx = _frame_filter_to_index(frame_filter, num_frames)
    frame_idx = max(0, min(num_frames - 1, frame_idx))
    depth_slice = depth_np[frame_idx]

    finite_mask = np.isfinite(depth_slice)
    if not finite_mask.any():
        return None

    valid_vals = depth_slice[finite_mask]
    d_min = float(valid_vals.min())
    d_max = float(valid_vals.max())
    if d_max - d_min < 1e-6:
        norm = np.zeros_like(depth_slice, dtype=np.float32)
    else:
        norm = (depth_slice - d_min) / (d_max - d_min)
        norm = np.clip(norm, 0.0, 1.0)

    norm_uint8 = (norm * 255.0).astype(np.uint8)
    colored = cv2.applyColorMap(norm_uint8, cv2.COLORMAP_PLASMA)
    colored = cv2.cvtColor(colored, cv2.COLOR_BGR2RGB)
    return colored


def build_pointcloud_from_predictions(
    predictions: dict,
    conf_thres: float = 50.0,
    frame_filter: str | None = "All",
    mask_black_bg: bool = False,
    mask_white_bg: bool = False,
):
    """
    Build a point cloud (Nx3) and optional RGB colors (Nx3, uint8) from predictions.

    - Uses predictions["world_points_from_depth"] (S,H,W,3)
    - Filters by predictions["depth_conf"] (S,H,W,1) with given conf_thres (0-100)
    - Optionally removes near-black/near-white background using predictions["images"] if present
    - Optionally selects a single frame via frame_filter like "3: filename"
    """
    wp = predictions.get("world_points_from_depth")
    if wp is None:
        raise ValueError("Predictions missing 'world_points_from_depth'; cannot export point cloud.")

    depth_conf = predictions.get("depth_conf")
    images = predictions.get("images")  # optional, for color

    # Select frames
    sel = _parse_frame_filter(frame_filter)
    pts = wp[sel]
    conf = depth_conf[sel] if depth_conf is not None else None
    img = images[sel] if images is not None else None

    pts = np.asarray(pts)
    if pts.ndim != 4 or pts.shape[-1] != 3:
        raise ValueError(f"Unexpected world_points_from_depth shape: {pts.shape}")

    S, H, W, _ = pts.shape
    pts_flat = pts.reshape(-1, 3)

    # Base validity mask
    mask = np.isfinite(pts_flat).all(axis=1)

    # Confidence filtering
    if conf is not None:
        conf_flat = np.asarray(conf).reshape(-1)
        thr = float(conf_thres) / 100.0
        mask &= (conf_flat >= thr)

    colors_flat = None
    if img is not None:
        img = np.asarray(img)
        # Normalize images to uint8 if needed
        if img.dtype != np.uint8:
            img_norm = img.astype(np.float32)
            # Handle potential [-1,1] range
            if img_norm.min() < 0.0:
                img_norm = (img_norm + 1.0) / 2.0
            img_norm = np.clip(img_norm, 0.0, 1.0)
            img_u8 = (img_norm * 255.0).round().astype(np.uint8)
        else:
            img_u8 = img

        colors_flat = img_u8.reshape(-1, 3)

        # Background filtering based on color
        if mask_black_bg:
            mask &= (colors_flat.sum(axis=1) > 3)  # not near pure black
        if mask_white_bg:
            mask &= (colors_flat.sum(axis=1) < (3 * 255 - 3))  # not near pure white

        colors_flat = colors_flat[mask]

    pts_final = pts_flat[mask]
    return pts_final, colors_flat


def save_pointcloud_ply(points: np.ndarray, colors: np.ndarray | None, out_path: str):
    """Save point cloud to ASCII PLY. Points: (N,3) float. Colors: (N,3) uint8 or None."""
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    N = int(points.shape[0])
    has_color = colors is not None and colors.shape[0] == N

    header = [
        "ply",
        "format ascii 1.0",
        f"element vertex {N}",
        "property float x",
        "property float y",
        "property float z",
    ]
    if has_color:
        header += [
            "property uchar red",
            "property uchar green",
            "property uchar blue",
        ]
    header += ["end_header\n"]

    with open(out_path, "w", encoding="utf-8") as f:
        f.write("\n".join(header))
        if has_color:
            for p, c in zip(points, colors):
                f.write(f"{float(p[0]):.6f} {float(p[1]):.6f} {float(p[2]):.6f} {int(c[0])} {int(c[1])} {int(c[2])}\n")
        else:
            for p in points:
                f.write(f"{float(p[0]):.6f} {float(p[1]):.6f} {float(p[2]):.6f}\n")


def save_pointcloud_npz(points: np.ndarray, colors: np.ndarray | None, out_path: str):
    """Save point cloud arrays to NPZ (points float32, colors uint8 if present)."""
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    data = {"points": points.astype(np.float32)}
    if colors is not None:
        data["colors"] = colors.astype(np.uint8)
    np.savez_compressed(out_path, **data)

def load_ckpt_flex(model: torch.nn.Module, ckpt_path: str, map_location: str = "cpu") -> torch.nn.Module:
    """
    加载多种常见格式的 checkpoint：
    - 支持顶层是 {state_dict|model|student|model_ema|student_ema}
    - 自动去除常见前缀：module./model./student.
    - 过滤掉与当前模型 shape 不匹配的权重，避免加载报错
    - 打印加载统计：loaded/ignored/missing/unexpected
    """
    if not os.path.exists(ckpt_path):
        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")

    raw = torch.load(ckpt_path, map_location=map_location)

    # 提取真正的权重字典
    if isinstance(raw, dict):
        for k in ["student", "student_ema", "model", "model_ema", "state_dict"]:
            if k in raw and isinstance(raw[k], dict):
                raw = raw[k]
                break

    if not isinstance(raw, dict):
        raise ValueError(f"Unsupported checkpoint format: {type(raw)}")

    # 去前缀
    def strip_prefix(d: dict, prefixes=("module.", "model.", "student.")):
        out = {}
        for k, v in d.items():
            nk = k
            for p in prefixes:
                if nk.startswith(p):
                    nk = nk[len(p):]
            out[nk] = v
        return out

    state = strip_prefix(raw)

    # 过滤形状不匹配的键，避免 _IncompatibleKeys 报 shape mismatch
    model_sd = model.state_dict()
    loadable = {}
    dropped = []
    for k, v in state.items():
        if k in model_sd and hasattr(v, 'shape') and hasattr(model_sd[k], 'shape'):
            if tuple(v.shape) == tuple(model_sd[k].shape):
                loadable[k] = v
            else:
                dropped.append((k, tuple(v.shape), tuple(model_sd[k].shape)))
        elif k in model_sd:
            # 有些张量可能不是典型的 torch.Tensor，但通常不该出现，保守丢弃
            dropped.append((k, getattr(v, 'shape', None), getattr(model_sd[k], 'shape', None)))

    incompatible = model.load_state_dict(loadable, strict=False)

    print(f"[load_ckpt] loaded keys: {len(loadable)}, dropped (shape-mismatch): {len(dropped)}")
    if dropped:
        print("  dropped samples (up to 10):", dropped[:10])
    if getattr(incompatible, 'missing_keys', None):
        print(f"[load_ckpt] missing keys: {len(incompatible.missing_keys)} (showing up to 10)")
        print("  ", incompatible.missing_keys[:10])
    if getattr(incompatible, 'unexpected_keys', None):
        print(f"[load_ckpt] unexpected keys: {len(incompatible.unexpected_keys)} (showing up to 10)")
        print("  ", incompatible.unexpected_keys[:10])

    return model


# 解析命令行或环境变量传入的 ckpt 路径
DEFAULT_CKPT_PATH = "/data0/liqifeng/ZYC/model.pt"
DEFAULT_FEATURE_DIR_PATH = \
    "/data1/datasets/VGGT_features/libero_10/KITCHEN_SCENE4_put_the_black_bowl_in_the_bottom_drawer_of_the_cabinet_and_close_it/sample_00000/iter_0001"
FEATURE_BASE_DIR = "/data1/datasets/VGGT_features"

parser = argparse.ArgumentParser(add_help=False)
parser.add_argument(
    "--ckpt",
    type=str,
    default=os.environ.get("VGGT_CKPT", DEFAULT_CKPT_PATH),
    help="Path to student/teacher checkpoint (.pt/.ckpt)",
)
parser.add_argument(
    "--feature-dir",
    type=str,
    default=os.environ.get("VGGT_FEATURE_DIR", DEFAULT_FEATURE_DIR_PATH),
    help="Path to a directory containing cached aggregator features",
)
# 允许忽略其他未知参数（便于与 gradio 共存）
args, _unknown = parser.parse_known_args()

local_ckpt = args.ckpt or DEFAULT_CKPT_PATH
default_feature_dir = args.feature_dir or DEFAULT_FEATURE_DIR_PATH

# 构建与加载模型
model = VGGT()
model = load_ckpt_flex(model, local_ckpt, map_location=device)
model.eval()
model = model.to(device)


# -------------------------------------------------------------------------
# 1) Core model inference
# -------------------------------------------------------------------------
def run_model(target_dir, model) -> dict:
    """
    Run the VGGT model on images in the 'target_dir/images' folder and return predictions.
    """
    print(f"Processing images from {target_dir}")

    # Device check
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if not torch.cuda.is_available():
        raise ValueError("CUDA is not available. Check your environment.")

    # Move model to device
    model = model.to(device)
    model.eval()

    # Load and preprocess images
    image_names = glob.glob(os.path.join(target_dir, "images", "*"))
    image_names = sorted(image_names)
    print(f"Found {len(image_names)} images")
    if len(image_names) == 0:
        raise ValueError("No images found. Check your upload.")

    images = load_and_preprocess_images(image_names).to(device)
    print(f"Preprocessed images shape: {images.shape}")

    # Run inference
    print("Running inference...")
    dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16

    with torch.no_grad():
        with torch.cuda.amp.autocast(dtype=dtype):
            predictions = model(images)

    # Convert pose encoding to extrinsic and intrinsic matrices
    print("Converting pose encoding to extrinsic and intrinsic matrices...")
    extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:])
    predictions["extrinsic"] = extrinsic
    predictions["intrinsic"] = intrinsic

    # Convert tensors to numpy
    for key in list(predictions.keys()):
        if isinstance(predictions[key], torch.Tensor):
            predictions[key] = predictions[key].cpu().numpy().squeeze(0)  # remove batch dimension
    # 移除不需要/不可序列化的字段，避免 np.savez 报错
    predictions.pop('pose_enc_list', None)

    # Generate world points from depth map
    print("Computing world points from depth map...")
    depth_map = predictions["depth"]  # (S, H, W, 1)
    world_points = unproject_depth_map_to_point_map(depth_map, predictions["extrinsic"], predictions["intrinsic"])
    predictions["world_points_from_depth"] = world_points

    # Clean up
    torch.cuda.empty_cache()
    return predictions, image_names


def _decode_h5_strings(dataset) -> list[str]:
    values = dataset[()]
    if isinstance(values, (bytes, np.bytes_, str)):
        values = [values]
    else:
        values = values.tolist()

    decoded = []
    for item in values:
        if isinstance(item, (bytes, np.bytes_)):
            decoded.append(item.decode("utf-8"))
        else:
            decoded.append(str(item))
    return decoded


def _load_scene_hdf5_int8_bundle(metadata: dict, feature_dir: str, device: str):
    if h5py is None:
        raise RuntimeError("h5py is required to read consolidated scene feature caches. Please install h5py.")

    scene_file_ref = metadata.get("scene_file")
    if not scene_file_ref:
        raise KeyError("Cached metadata missing 'scene_file' for scene-level feature bundle")

    if os.path.isabs(scene_file_ref):
        scene_file_path = scene_file_ref
    else:
        scene_file_path = os.path.normpath(os.path.join(feature_dir, scene_file_ref))

    if not os.path.exists(scene_file_path):
        raise FileNotFoundError(f"Scene HDF5 file not found: {scene_file_path}")

    sample_group = metadata.get("sample_group")
    iteration_group = metadata.get("iteration_group")
    if not sample_group or not iteration_group:
        raise KeyError("Cached metadata missing 'sample_group' or 'iteration_group'")

    saved_layers = metadata.get("saved_layers", [])
    layer_arrays: dict[int, np.ndarray] = {}
    image_paths = metadata.get("image_paths", [])

    with h5py.File(scene_file_path, "r") as scene_file:
        if sample_group not in scene_file:
            raise KeyError(f"Sample group '{sample_group}' not present in {scene_file_path}")
        sample_grp = scene_file[sample_group]
        if iteration_group not in sample_grp:
            raise KeyError(f"Iteration group '{iteration_group}' not present in scene HDF5")
        iter_grp = sample_grp[iteration_group]

        quant_features = iter_grp["features"][...]
        feature_scale = float(iter_grp.attrs.get("feature_scale", metadata.get("feature_scale", 1.0)))
        final_array = quant_features.astype(np.float32) * feature_scale

        layer_grp = iter_grp.get("layer_features")
        layer_scales_meta = metadata.get("layer_scales", {})
        for layer_name in saved_layers:
            if layer_grp is None or layer_name not in layer_grp:
                continue
            dataset = layer_grp[layer_name]
            layer_scale = float(dataset.attrs.get("scale", layer_scales_meta.get(layer_name, feature_scale)))
            layer_array = dataset[...].astype(np.float32) * layer_scale
            try:
                layer_idx = int(layer_name.split("_")[1])
            except (IndexError, ValueError):
                continue
            layer_arrays[layer_idx] = layer_array

        if not image_paths and "image_paths" in iter_grp:
            image_paths = _decode_h5_strings(iter_grp["image_paths"])

    final_tensor = torch.from_numpy(final_array.astype(np.float32)).to(device=device, dtype=torch.float32)
    layer_tensors = {
        idx: torch.from_numpy(arr.astype(np.float32)).to(device=device, dtype=torch.float32)
        for idx, arr in layer_arrays.items()
    }

    max_saved_idx = max(layer_tensors.keys(), default=-1)
    total_layers = metadata.get("num_aggregator_iterations")
    if total_layers is None:
        total_layers = max_saved_idx + 1 if max_saved_idx >= 0 else 1

    aggregated_tokens_list = []
    for idx in range(total_layers):
        tensor = layer_tensors.get(idx, final_tensor)
        aggregated_tokens_list.append(tensor)

    if aggregated_tokens_list:
        aggregated_tokens_list[-1] = final_tensor
    else:
        aggregated_tokens_list.append(final_tensor)

    preprocess_mode = metadata.get("preprocess_mode") or "crop"
    patch_start_idx = metadata["patch_start_idx"]

    return aggregated_tokens_list, patch_start_idx, image_paths, preprocess_mode


def load_cached_feature_bundle(feature_dir: str, device: str):
    metadata_path = os.path.join(feature_dir, "metadata.json")
    if not os.path.exists(metadata_path):
        raise FileNotFoundError(f"Metadata file not found: {metadata_path}")

    with open(metadata_path, "r", encoding="utf-8") as f:
        metadata = json.load(f)

    if "patch_start_idx" not in metadata:
        raise KeyError("Cached metadata missing 'patch_start_idx'")

    if metadata.get("storage_format") == "scene_hdf5_int8":
        return _load_scene_hdf5_int8_bundle(metadata, feature_dir, device)

    patch_start_idx = metadata["patch_start_idx"]
    preprocess_mode = metadata.get("preprocess_mode") or "crop"
    saved_layers = metadata.get("saved_layers", [])
    total_layers = metadata.get("num_aggregator_iterations")

    layer_tensors = {}
    max_saved_idx = -1
    for layer_name in saved_layers:
        try:
            idx = int(layer_name.split("_")[1])
        except (IndexError, ValueError):
            continue
        layer_path_h5 = os.path.join(feature_dir, f"aggregator_features_{layer_name}.h5")
        layer_path_npz = os.path.join(feature_dir, f"aggregator_features_{layer_name}.npz")
        layer_path_npy = os.path.join(feature_dir, f"aggregator_features_{layer_name}.npy")

        if os.path.exists(layer_path_h5):
            if h5py is None:
                raise RuntimeError(
                    "h5py is required to read HDF5 feature caches. Please install h5py to continue."
                )
            with h5py.File(layer_path_h5, "r") as f:
                layer_array = f["features"][:]
        elif os.path.exists(layer_path_npz):
            with np.load(layer_path_npz, allow_pickle=False) as data:
                layer_array = data["features"]
        elif os.path.exists(layer_path_npy):
            layer_array = np.load(layer_path_npy, allow_pickle=False)
        else:
            continue
        layer_tensors[idx] = torch.from_numpy(layer_array).to(device=device, dtype=torch.float32)
        max_saved_idx = max(max_saved_idx, idx)

    final_path_h5 = os.path.join(feature_dir, "aggregator_features.h5")
    final_path_npz = os.path.join(feature_dir, "aggregator_features.npz")
    final_path_npy = os.path.join(feature_dir, "aggregator_features.npy")

    if os.path.exists(final_path_h5):
        if h5py is None:
            raise RuntimeError("h5py is required to read HDF5 feature caches. Please install h5py to continue.")
        with h5py.File(final_path_h5, "r") as f:
            final_array = f["features"][:]
    elif os.path.exists(final_path_npz):
        with np.load(final_path_npz, allow_pickle=False) as data:
            final_array = data["features"]
    elif os.path.exists(final_path_npy):
        final_array = np.load(final_path_npy, allow_pickle=False)
    else:
        raise FileNotFoundError(
            f"Cached features missing file: {final_path_h5} (or legacy {final_path_npz}/{final_path_npy})"
        )
    final_tensor = torch.from_numpy(final_array).to(device=device, dtype=torch.float32)

    if total_layers is None:
        total_layers = max(max_saved_idx, 0) + 1

    aggregated_tokens_list = []
    for idx in range(total_layers):
        tensor = layer_tensors.get(idx)
        if tensor is None:
            tensor = final_tensor
        aggregated_tokens_list.append(tensor)

    if aggregated_tokens_list:
        aggregated_tokens_list[-1] = final_tensor
    else:
        aggregated_tokens_list.append(final_tensor)

    image_paths = metadata.get("image_paths")
    if image_paths is None:
        image_paths_path = os.path.join(feature_dir, "image_paths.json")
        if not os.path.exists(image_paths_path):
            raise FileNotFoundError(f"Cached features missing image list: {image_paths_path}")
        with open(image_paths_path, "r", encoding="utf-8") as f:
            image_paths = json.load(f)

    return aggregated_tokens_list, patch_start_idx, image_paths, preprocess_mode


def run_model_from_cached_features(feature_dir: str, model) -> tuple[dict, list[str]]:
    print(f"Loading cached features from {feature_dir}")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device != "cuda":
        raise ValueError("CUDA is required for running VGGT heads")

    aggregated_tokens_list, patch_start_idx, image_paths, preprocess_mode = load_cached_feature_bundle(
        feature_dir, device
    )

    if len(image_paths) == 0:
        raise ValueError("Cached feature directory did not include any image paths")

    images = load_and_preprocess_images(image_paths, mode=preprocess_mode).to(device)
    images_batched = images.unsqueeze(0)

    model = model.to(device)
    model.eval()

    predictions = {}

    with torch.no_grad():
        with torch.cuda.amp.autocast(enabled=False):
            if model.camera_head is not None:
                pose_enc_list = model.camera_head(aggregated_tokens_list)
                predictions["pose_enc"] = pose_enc_list[-1]
                predictions["pose_enc_list"] = pose_enc_list

            if model.depth_head is not None:
                depth, depth_conf = model.depth_head(
                    aggregated_tokens_list, images=images_batched, patch_start_idx=patch_start_idx
                )
                predictions["depth"] = depth
                predictions["depth_conf"] = depth_conf

            if model.point_head is not None:
                pts3d, pts3d_conf = model.point_head(
                    aggregated_tokens_list, images=images_batched, patch_start_idx=patch_start_idx
                )
                predictions["world_points"] = pts3d
                predictions["world_points_conf"] = pts3d_conf

    if not model.training:
        predictions["images"] = images_batched

    extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images_batched.shape[-2:])
    predictions["extrinsic"] = extrinsic
    predictions["intrinsic"] = intrinsic

    for key in list(predictions.keys()):
        if isinstance(predictions[key], torch.Tensor):
            predictions[key] = predictions[key].detach().cpu().numpy().squeeze(0)

    predictions.pop("pose_enc_list", None)

    depth_map = predictions["depth"]
    world_points = unproject_depth_map_to_point_map(depth_map, predictions["extrinsic"], predictions["intrinsic"])
    predictions["world_points_from_depth"] = world_points

    torch.cuda.empty_cache()
    return predictions, image_paths


# -------------------------------------------------------------------------
# 2) Handle uploaded video/images --> produce target_dir + images
# -------------------------------------------------------------------------
def handle_uploads(input_video, input_images):
    """
    Create a new 'target_dir' + 'images' subfolder, and place user-uploaded
    images or extracted frames from video into it. Return (target_dir, image_paths).
    """
    start_time = time.time()
    gc.collect()
    torch.cuda.empty_cache()

    # Create a unique folder name
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    target_dir = f"input_images_{timestamp}"
    target_dir_images = os.path.join(target_dir, "images")

    # Clean up if somehow that folder already exists
    if os.path.exists(target_dir):
        shutil.rmtree(target_dir)
    os.makedirs(target_dir)
    os.makedirs(target_dir_images)

    image_paths = []

    # --- Handle images ---
    if input_images is not None:
        for file_data in input_images:
            if isinstance(file_data, dict) and "name" in file_data:
                file_path = file_data["name"]
            else:
                file_path = file_data
            dst_path = os.path.join(target_dir_images, os.path.basename(file_path))
            shutil.copy(file_path, dst_path)
            image_paths.append(dst_path)

    # --- Handle video ---
    if input_video is not None:
        if isinstance(input_video, dict) and "name" in input_video:
            video_path = input_video["name"]
        else:
            video_path = input_video

        vs = cv2.VideoCapture(video_path)
        fps = vs.get(cv2.CAP_PROP_FPS)
        frame_interval = int(fps * 1)  # 1 frame/sec

        count = 0
        video_frame_num = 0
        while True:
            gotit, frame = vs.read()
            if not gotit:
                break
            count += 1
            if count % frame_interval == 0:
                image_path = os.path.join(target_dir_images, f"{video_frame_num:06}.png")
                cv2.imwrite(image_path, frame)
                image_paths.append(image_path)
                video_frame_num += 1

    # Sort final images for gallery
    image_paths = sorted(image_paths)

    end_time = time.time()
    print(f"Files copied to {target_dir_images}; took {end_time - start_time:.3f} seconds")
    return target_dir, image_paths


# -------------------------------------------------------------------------
# 3) Update gallery on upload
# -------------------------------------------------------------------------
def update_gallery_on_upload(input_video, input_images):
    """
    Whenever user uploads or changes files, immediately handle them
    and show in the gallery. Return (target_dir, image_paths).
    If nothing is uploaded, returns "None" and empty list.
    """
    if not input_video and not input_images:
        return None, None, None, None
    target_dir, image_paths = handle_uploads(input_video, input_images)
    return None, target_dir, image_paths, "Upload complete. Click 'Reconstruct' to begin 3D processing."


# -------------------------------------------------------------------------
# 4) Reconstruction: uses the target_dir plus any viz parameters
# -------------------------------------------------------------------------
def gradio_demo(
    target_dir,
    conf_thres=3.0,
    frame_filter="All",
    mask_black_bg=False,
    mask_white_bg=False,
    show_cam=True,
    mask_sky=False,
    prediction_mode="Pointmap Regression",
    cached_feature_dir=None,
):
    """
    Perform reconstruction using the already-created target_dir/images.
    """
    use_cached = cached_feature_dir and os.path.isdir(cached_feature_dir)

    if not use_cached and (not os.path.isdir(target_dir) or target_dir == "None"):
        return None, "No valid target directory found. Please upload first.", None, target_dir

    current_dir = cached_feature_dir if use_cached else target_dir

    start_time = time.time()
    gc.collect()
    torch.cuda.empty_cache()

    if use_cached:
        frame_filter_choices = None
    else:
        target_dir_images = os.path.join(current_dir, "images")
        all_files = sorted(os.listdir(target_dir_images)) if os.path.isdir(target_dir_images) else []
        frame_filter_choices = ["All"] + [f"{i}: {filename}" for i, filename in enumerate(all_files)]

    print("Running VGGT heads...")
    with torch.no_grad():
        if use_cached:
            predictions, image_paths = run_model_from_cached_features(current_dir, model)
            frame_filter_choices = ["All"] + [f"{i}: {os.path.basename(p)}" for i, p in enumerate(image_paths)]
        else:
            predictions, image_paths = run_model(current_dir, model)
            if frame_filter_choices is None:
                frame_filter_choices = ["All"] + [f"{i}: {os.path.basename(p)}" for i, p in enumerate(image_paths)]

    # Save predictions
    prediction_save_path = os.path.join(current_dir, "predictions.npz")
    np.savez(prediction_save_path, **predictions)

    # Handle None frame_filter
    if frame_filter is None:
        frame_filter = "All"

    depth_preview = _render_depth_preview(predictions.get("depth"), frame_filter)

    # Build a GLB file name
    glbfile = os.path.join(
        current_dir,
        f"glbscene_{conf_thres}_{frame_filter.replace('.', '_').replace(':', '').replace(' ', '_')}_maskb{mask_black_bg}_maskw{mask_white_bg}_cam{show_cam}_sky{mask_sky}_pred{prediction_mode.replace(' ', '_')}.glb",
    )

    # Convert predictions to GLB
    glbscene = predictions_to_glb(
        predictions,
        conf_thres=conf_thres,
        filter_by_frames=frame_filter,
        mask_black_bg=mask_black_bg,
        mask_white_bg=mask_white_bg,
        show_cam=show_cam,
        mask_sky=mask_sky,
        target_dir=current_dir,
        prediction_mode=prediction_mode,
    )
    glbscene.export(file_obj=glbfile)

    # Additionally, export point cloud to PLY and NPZ for external usage
    try:
        pts, cols = build_pointcloud_from_predictions(
            predictions,
            conf_thres=conf_thres,
            frame_filter=frame_filter,
            mask_black_bg=mask_black_bg,
            mask_white_bg=mask_white_bg,
        )

        pc_suffix = f"{conf_thres}_{frame_filter.replace('.', '_').replace(':', '').replace(' ', '_')}_maskb{mask_black_bg}_maskw{mask_white_bg}"
        plyfile = os.path.join(current_dir, f"pointcloud_{pc_suffix}.ply")
        npzfile = os.path.join(current_dir, f"pointcloud_{pc_suffix}.npz")
        save_pointcloud_ply(pts, cols, plyfile)
        save_pointcloud_npz(pts, cols, npzfile)
        pc_log = f" Point cloud saved: {os.path.basename(plyfile)} and {os.path.basename(npzfile)}."
    except Exception as e:
        pc_log = f" Point cloud export skipped: {e}."

    # Cleanup
    del predictions
    gc.collect()
    torch.cuda.empty_cache()

    end_time = time.time()
    print(f"Total time: {end_time - start_time:.2f} seconds (including IO)")
    frame_count = len(image_paths)
    log_msg = f"Reconstruction Success ({frame_count} frames). Waiting for visualization." + pc_log

    dropdown = gr.Dropdown(choices=frame_filter_choices, value=frame_filter, interactive=True)

    return glbfile, depth_preview, log_msg, dropdown, current_dir


# -------------------------------------------------------------------------
# 5) Helper functions for UI resets + re-visualization
# -------------------------------------------------------------------------
def clear_fields():
    """
    Clears the 3D viewer, the stored target_dir, and empties the gallery.
    """
    return None, None


def update_log():
    """
    Display a quick log message while waiting.
    """
    return "Loading and Reconstructing..."


def update_visualization(
    target_dir, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example
):
    """
    Reload saved predictions from npz, create (or reuse) the GLB for new parameters,
    and return it for the 3D viewer. If is_example == "True", skip.
    """

    # If it's an example click, skip as requested
    if is_example == "True":
        return None, "No reconstruction available. Please click the Reconstruct button first.", None

    if not target_dir or target_dir == "None" or not os.path.isdir(target_dir):
        return None, "No reconstruction available. Please click the Reconstruct button first.", None

    predictions_path = os.path.join(target_dir, "predictions.npz")
    if not os.path.exists(predictions_path):
        return None, f"No reconstruction available at {predictions_path}. Please run 'Reconstruct' first.", None

    key_list = [
        "pose_enc",
        "depth",
        "depth_conf",
        "world_points",
        "world_points_conf",
        "images",
        "extrinsic",
        "intrinsic",
        "world_points_from_depth",
    ]

    loaded = np.load(predictions_path)
    predictions = {key: np.array(loaded[key]) for key in key_list}

    glbfile = os.path.join(
        target_dir,
        f"glbscene_{conf_thres}_{frame_filter.replace('.', '_').replace(':', '').replace(' ', '_')}_maskb{mask_black_bg}_maskw{mask_white_bg}_cam{show_cam}_sky{mask_sky}_pred{prediction_mode.replace(' ', '_')}.glb",
    )

    if not os.path.exists(glbfile):
        glbscene = predictions_to_glb(
            predictions,
            conf_thres=conf_thres,
            filter_by_frames=frame_filter,
            mask_black_bg=mask_black_bg,
            mask_white_bg=mask_white_bg,
            show_cam=show_cam,
            mask_sky=mask_sky,
            target_dir=target_dir,
            prediction_mode=prediction_mode,
        )
        glbscene.export(file_obj=glbfile)

    depth_preview = _render_depth_preview(predictions.get("depth"), frame_filter)

    return glbfile, "Updating Visualization", depth_preview


# -------------------------------------------------------------------------
# Example images
# -------------------------------------------------------------------------

great_wall_video = "examples/videos/great_wall.mp4"
colosseum_video = "examples/videos/Colosseum.mp4"
room_video = "examples/videos/room.mp4"
kitchen_video = "examples/videos/kitchen.mp4"
fern_video = "examples/videos/fern.mp4"
single_cartoon_video = "examples/videos/single_cartoon.mp4"
single_oil_painting_video = "examples/videos/single_oil_painting.mp4"
pyramid_video = "examples/videos/pyramid.mp4"


# -------------------------------------------------------------------------
# 6) Build Gradio UI
# -------------------------------------------------------------------------
theme = gr.themes.Ocean()
theme.set(
    checkbox_label_background_fill_selected="*button_primary_background_fill",
    checkbox_label_text_color_selected="*button_primary_text_color",
)

with gr.Blocks(
    theme=theme,
    css="""
    .custom-log * {
        font-style: italic;
        font-size: 22px !important;
        background-image: linear-gradient(120deg, #0ea5e9 0%, #6ee7b7 60%, #34d399 100%);
        -webkit-background-clip: text;
        background-clip: text;
        font-weight: bold !important;
        color: transparent !important;
        text-align: center !important;
    }
    
    .example-log * {
        font-style: italic;
        font-size: 16px !important;
        background-image: linear-gradient(120deg, #0ea5e9 0%, #6ee7b7 60%, #34d399 100%);
        -webkit-background-clip: text;
        background-clip: text;
        color: transparent !important;
    }
    
    #my_radio .wrap {
        display: flex;
        flex-wrap: nowrap;
        justify-content: center;
        align-items: center;
    }

    #my_radio .wrap label {
        display: flex;
        width: 50%;
        justify-content: center;
        align-items: center;
        margin: 0;
        padding: 10px 0;
        box-sizing: border-box;
    }
    """,
) as demo:
    # Instead of gr.State, we use a hidden Textbox:
    is_example = gr.Textbox(label="is_example", visible=False, value="None")
    num_images = gr.Textbox(label="num_images", visible=False, value="None")

    gr.HTML(
        """
    <h1>🏛️ VGGT: Visual Geometry Grounded Transformer</h1>
    <p>
    <a href="https://github.com/facebookresearch/vggt">🐙 GitHub Repository</a> |
    <a href="#">Project Page</a>
    </p>

    <div style="font-size: 16px; line-height: 1.5;">
    <p>Upload a video or a set of images to create a 3D reconstruction of a scene or object. VGGT takes these images and generates a 3D point cloud, along with estimated camera poses.</p>

    <h3>Getting Started:</h3>
    <ol>
        <li><strong>Upload Your Data:</strong> Use the "Upload Video" or "Upload Images" buttons on the left to provide your input. Videos will be automatically split into individual frames (one frame per second).</li>
        <li><strong>Preview:</strong> Your uploaded images will appear in the gallery on the left.</li>
        <li><strong>Reconstruct:</strong> Click the "Reconstruct" button to start the 3D reconstruction process.</li>
        <li><strong>Visualize:</strong> The 3D reconstruction will appear in the viewer on the right. You can rotate, pan, and zoom to explore the model, and download the GLB file. Note the visualization of 3D points may be slow for a large number of input images.</li>
        <li>
        <strong>Adjust Visualization (Optional):</strong>
        After reconstruction, you can fine-tune the visualization using the options below
        <details style="display:inline;">
            <summary style="display:inline;">(<strong>click to expand</strong>):</summary>
            <ul>
            <li><em>Confidence Threshold:</em> Adjust the filtering of points based on confidence.</li>
            <li><em>Show Points from Frame:</em> Select specific frames to display in the point cloud.</li>
            <li><em>Show Camera:</em> Toggle the display of estimated camera positions.</li>
            <li><em>Filter Sky / Filter Black Background:</em> Remove sky or black-background points.</li>
            <li><em>Select a Prediction Mode:</em> Choose between "Depthmap and Camera Branch" or "Pointmap Branch."</li>
            </ul>
        </details>
        </li>
    </ol>
    <p><strong style="color: #0ea5e9;">Please note:</strong> <span style="color: #0ea5e9; font-weight: bold;">VGGT typically reconstructs a scene in less than 1 second. However, visualizing 3D points may take tens of seconds due to third-party rendering, which are independent of VGGT's processing time. </span></p>
    </div>
    """
    )

    target_dir_output = gr.Textbox(label="Target Dir", visible=False, value="None")

    with gr.Row():
        with gr.Column(scale=2):
            input_video = gr.Video(label="Upload Video", interactive=True)
            input_images = gr.File(file_count="multiple", label="Upload Images", interactive=True)

            feature_dir_input = gr.Textbox(
                label="Cached Feature Directory",
                value=default_feature_dir or "",
                placeholder="Optional: path to precomputed feature directory",
            )

            image_gallery = gr.Gallery(
                label="Preview",
                columns=4,
                height="300px",
                show_download_button=True,
                object_fit="contain",
                preview=True,
            )

        with gr.Column(scale=4):
            with gr.Column():
                gr.Markdown("**3D Reconstruction (Point Cloud and Camera Poses)**")
                log_output = gr.Markdown(
                    "Please upload a video or images, then click Reconstruct.", elem_classes=["custom-log"]
                )
                reconstruction_output = gr.Model3D(height=520, zoom_speed=0.5, pan_speed=0.5)
                depth_preview = gr.Image(label="Depth Preview", height=260)

            with gr.Row():
                submit_btn = gr.Button("Reconstruct", scale=1, variant="primary")
                clear_btn = gr.ClearButton(
                    [
                        input_video,
                        input_images,
                        reconstruction_output,
                        depth_preview,
                        log_output,
                        target_dir_output,
                        image_gallery,
                    ],
                    scale=1,
                )

            with gr.Row():
                prediction_mode = gr.Radio(
                    ["Depthmap and Camera Branch", "Pointmap Branch"],
                    label="Select a Prediction Mode",
                    value="Depthmap and Camera Branch",
                    scale=1,
                    elem_id="my_radio",
                )

            with gr.Row():
                conf_thres = gr.Slider(minimum=0, maximum=100, value=50, step=0.1, label="Confidence Threshold (%)")
                frame_filter = gr.Dropdown(choices=["All"], value="All", label="Show Points from Frame")
                with gr.Column():
                    show_cam = gr.Checkbox(label="Show Camera", value=True)
                    mask_sky = gr.Checkbox(label="Filter Sky", value=False)
                    mask_black_bg = gr.Checkbox(label="Filter Black Background", value=False)
                    mask_white_bg = gr.Checkbox(label="Filter White Background", value=False)

    # ---------------------- Examples section ----------------------
    examples = [
        [colosseum_video, "22", None, 20.0, False, False, True, False, "Depthmap and Camera Branch", "True"],
        [pyramid_video, "30", None, 35.0, False, False, True, False, "Depthmap and Camera Branch", "True"],
        [single_cartoon_video, "1", None, 15.0, False, False, True, False, "Depthmap and Camera Branch", "True"],
        [single_oil_painting_video, "1", None, 20.0, False, False, True, True, "Depthmap and Camera Branch", "True"],
        [room_video, "8", None, 5.0, False, False, True, False, "Depthmap and Camera Branch", "True"],
        [kitchen_video, "25", None, 50.0, False, False, True, False, "Depthmap and Camera Branch", "True"],
        [fern_video, "20", None, 45.0, False, False, True, False, "Depthmap and Camera Branch", "True"],
    ]

    def example_pipeline(
        input_video,
        num_images_str,
        input_images,
        conf_thres,
        mask_black_bg,
        mask_white_bg,
        show_cam,
        mask_sky,
        prediction_mode,
        is_example_str,
    ):
        """
        1) Copy example images to new target_dir
        2) Reconstruct
        3) Return model3D + logs + new_dir + updated dropdown + gallery
        We do NOT return is_example. It's just an input.
        """
        target_dir, image_paths = handle_uploads(input_video, input_images)
        # Always use "All" for frame_filter in examples
        frame_filter = "All"
        glbfile, depth_image, log_msg, dropdown, used_dir = gradio_demo(
            target_dir,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            cached_feature_dir=None,
        )
        return glbfile, depth_image, log_msg, used_dir, dropdown, image_paths

    gr.Markdown("Click any row to load an example.", elem_classes=["example-log"])

    gr.Examples(
        examples=examples,
        inputs=[
            input_video,
            num_images,
            input_images,
            conf_thres,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            is_example,
        ],
        outputs=[reconstruction_output, depth_preview, log_output, target_dir_output, frame_filter, image_gallery],
        fn=example_pipeline,
        cache_examples=False,
        examples_per_page=50,
    )

    # -------------------------------------------------------------------------
    # "Reconstruct" button logic:
    #  - Clear fields
    #  - Update log
    #  - gradio_demo(...) with the existing target_dir
    #  - Then set is_example = "False"
    # -------------------------------------------------------------------------
    submit_btn.click(fn=clear_fields, inputs=[], outputs=[reconstruction_output, depth_preview]).then(
        fn=update_log, inputs=[], outputs=[log_output]
    ).then(
        fn=gradio_demo,
        inputs=[
            target_dir_output,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            feature_dir_input,
        ],
        outputs=[reconstruction_output, depth_preview, log_output, frame_filter, target_dir_output],
    ).then(
        fn=lambda: "False", inputs=[], outputs=[is_example]  # set is_example to "False"
    )

    # -------------------------------------------------------------------------
    # Real-time Visualization Updates
    # -------------------------------------------------------------------------
    conf_thres.change(
        update_visualization,
        [
            target_dir_output,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            is_example,
        ],
        [reconstruction_output, log_output, depth_preview],
    )
    frame_filter.change(
        update_visualization,
        [
            target_dir_output,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            is_example,
        ],
        [reconstruction_output, log_output, depth_preview],
    )
    mask_black_bg.change(
        update_visualization,
        [
            target_dir_output,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            is_example,
        ],
        [reconstruction_output, log_output, depth_preview],
    )
    mask_white_bg.change(
        update_visualization,
        [
            target_dir_output,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            is_example,
        ],
        [reconstruction_output, log_output, depth_preview],
    )
    show_cam.change(
        update_visualization,
        [
            target_dir_output,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            is_example,
        ],
        [reconstruction_output, log_output, depth_preview],
    )
    mask_sky.change(
        update_visualization,
        [
            target_dir_output,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            is_example,
        ],
        [reconstruction_output, log_output, depth_preview],
    )
    prediction_mode.change(
        update_visualization,
        [
            target_dir_output,
            conf_thres,
            frame_filter,
            mask_black_bg,
            mask_white_bg,
            show_cam,
            mask_sky,
            prediction_mode,
            is_example,
        ],
        [reconstruction_output, log_output, depth_preview],
    )

    # -------------------------------------------------------------------------
    # Auto-update gallery whenever user uploads or changes their files
    # -------------------------------------------------------------------------
    input_video.change(
        fn=update_gallery_on_upload,
        inputs=[input_video, input_images],
        outputs=[reconstruction_output, target_dir_output, image_gallery, log_output],
    )
    input_images.change(
        fn=update_gallery_on_upload,
        inputs=[input_video, input_images],
        outputs=[reconstruction_output, target_dir_output, image_gallery, log_output],
    )

    demo.queue(max_size=20).launch(
        show_error=True,
        share=False,
        server_name="0.0.0.0",
        server_port=8766,
        allowed_paths=[FEATURE_BASE_DIR, DEFAULT_FEATURE_DIR_PATH],
    )
